from typing import List, Optional, Any
from FlagEmbedding import LayerWiseFlagLLMReranker
from langchain_core.documents import Document
from pydantic import Field, PrivateAttr


class FlagLLMReranker:
    """Flag LLM Reranker."""
    model: str = Field(description="BAAI Reranker model name.")
    # device: int = Field(description="Device to use for inference.")
    _model: Any = PrivateAttr()

    def __init__(
        self,
        top_n: int = 5,
        model: str = "BAAI/bge-reranker-large",
        use_fp16: bool = True,
        cutoff_layers: List[int] = None,
        device: int = 0
    ) -> None:
        self._model = LayerWiseFlagLLMReranker(
            model,
            use_fp16=use_fp16,
            device=device
        )

        self.cutoff_layers = cutoff_layers if cutoff_layers is not None else [28]
        self.top_n = top_n


    @classmethod
    def class_name(cls) -> str:
        return "FlagLLMReranker"

    def postprocess_nodes(
        self,
        nodes: List[Document],
        query: str,
        limit: Optional[float] = 0,
    ) -> List[Document]:
        if len(nodes) == 0:
            return []

        new_nodes = []
        for node in nodes:
            scores = self._model.compute_score([[query, node.page_content]], cutoff_layers=self.cutoff_layers, normalize=True)
            score = scores[0][0]

            node.metadata["score"] = score

            if score > limit:
                new_nodes.append(node)

        new_nodes = sorted(new_nodes, key=lambda x: -x.metadata["score"] if x.metadata["score"] else 0)[: self.top_n]

        # 添加引用
        urls = {}
        for i in range(len(new_nodes)):
            if new_nodes[i].metadata["url"] == " ":
                new_nodes[i].metadata["url"] = new_nodes[i].metadata["title"]

            if new_nodes[i].metadata["url"] not in urls:
                urls[new_nodes[i].metadata["url"]] = len(urls) + 1
            new_nodes[i].page_content = f"[[citation:{urls[new_nodes[i].metadata['url']]}]] {new_nodes[i].page_content}"

        return new_nodes

